import numpy as np
import pickle
import torch
import torch.utils.data as torch_data

depth_per_simulation = 1
temperature = 1
Cpuct = 0.1
batch_size = 30
board_size = 15
learning_rate = 0.1
point = "tuple[int, int]"


class distribution_calculater:
    def __init__(self, size):
        self.map = {}
        self.order = []
        for i in range(size):
            for j in range(size):
                name = str(i) + "," + str(j)
                self.order.append(name)
                self.map[name] = 0

    def push(self, key, value):
        self.map[key] = value

    def get(self, train=True):
        result = []
        choice_pool = []
        choice_prob = []
        for key in self.order:
            if self.map[key] != 0:
                choice_pool.append(key)
                tmp = np.float_power(self.map[key], 1 / temperature)
                choice_prob.append(tmp)
                result.append(tmp)
                self.map[key] = 0
            else:
                result.append(0)

        he = sum(result)
        for i in range(len(result)):
            if result[i]:
                result[i] = result[i] / he
        choice_prob = [choice / he for choice in choice_prob]
        if train:
            move = np.random.choice(choice_pool, p=0.8 * np.array(choice_prob) + 0.2 * np.random.dirichlet(0.3 * np.ones(len(choice_prob))))
        else:
            move = choice_pool[np.argmax(choice_prob)]
        return move, result


def step_child_remove(board_pool, child_pool):
    i = 0
    while i < len(board_pool) and len(child_pool) != 0:
        j = 0
        while j < len(child_pool):
            if np.array_equal(board_pool[i], child_pool[j]):
                board_pool.pop(i)
                child_pool.pop(j)
                i -= 1
                break
            else:
                j += 1
        i += 1
    return board_pool


def write_file(object, file_name):
    filewriter = open(file_name, "wb")
    pickle.dump(object, filewriter)
    filewriter.close()


def read_file(file_name):
    filereader = open(file_name, "rb")
    object = pickle.load(filereader)
    filereader.close()
    return object


def move_to_str(action):
    return str(action[0]) + "," + str(action[1])


def str_to_move(string):
    move = string.split(",")
    return np.array([int(move[0]), int(move[1])])


def valid_move(state):
    return list(np.argwhere(state == 0))


def find_index(ndarray_list: list, target: np.ndarray):
    for i in range(len(ndarray_list)):
        if np.array_equal(ndarray_list[i], target):
            return i
    return None


class random_stack:
    def __init__(self, length=1000):
        self.state = []
        self.distrib = []
        self.winner = []
        self.length = length

    def isEmpty(self):
        return len(self.state) == 0

    def push(self, item):
        self.state.append(item["state"])
        self.distrib.append(item["distribution"])
        self.winner.append(item["value"])
        if len(self.state) >= self.length:
            self.state = self.state[1:]
            self.distrib = self.distrib[1:]
            self.winner = self.winner[1:]

    def seq(self):
        return self.state, self.distrib, self.winner


def generate_training_data(game_record, board_size):
    board = np.zeros([board_size, board_size])
    data = []
    player = 1
    winner = -1 if len(game_record) % 2 == 0 else 1
    for i in range(len(game_record)):
        step = str_to_move(game_record[i]["action"])
        state = transfer_to_input(board, player, board_size)
        data.append({"state": state, "distribution": game_record[i]["distribution"], "value": winner})
        board[step[0], step[1]] = player
        player, winner = -player, -winner
    return data


def generate_data_loader(stack):
    state, distrib, winner = stack.seq()
    tensor_x = torch.stack(tuple([torch.from_numpy(s) for s in state]))
    tensor_y1 = torch.stack(tuple([torch.Tensor(y1) for y1 in distrib]))
    tensor_y2 = torch.stack(tuple([torch.Tensor([float(y2)]) for y2 in winner]))
    dataset = torch_data.TensorDataset(tensor_x, tensor_y1, tensor_y2)
    my_loader = torch_data.DataLoader(dataset, batch_size=batch_size, shuffle=True)
    return my_loader


def transfer_to_input(state, current_player, board_size, history: "list[point]"):
    """
    history:list,length 2
    history[0]：上一步棋的位置，tuple
    history[1]：倒数第二步下棋的位置，tuple
    """
    # print(history)
    zero_board = np.zeros([board_size, board_size])
    if len(history) == 2:
        tmp1 = np.array(state.copy()).astype(float)
        state[history[1][0]][history[1][1]] = 0
        tmp2 = np.array(state.copy()).astype(float)
        state[history[0][0]][history[0][1]] = 0
        tmp3 = np.array(state.copy()).astype(float)
    elif len(history) == 1:
        tmp1 = tmp1 = np.array(state.copy()).astype(float)
        tmp2 = zero_board.copy()
        tmp3 = zero_board.copy()
    else:
        tmp1 = zero_board.copy()
        tmp2 = zero_board.copy()
        tmp3 = zero_board.copy()
    return np.stack([tmp1, tmp2, tmp3])
